{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "e:\\env\\anaconda3\\Lib\\site-packages\\transformers\\utils\\hub.py:106: FutureWarning: Using `TRANSFORMERS_CACHE` is deprecated and will be removed in v5 of Transformers. Use `HF_HOME` instead.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import json\n",
    "from tqdm import tqdm\n",
    "from loguru import logger\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import f1_score, classification_report, accuracy_score\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch import optim\n",
    "import torch\n",
    "\n",
    "class Config:\n",
    "    train_data = 'D:/project/python/大洋测试/data/train.txt'\n",
    "    test_data = 'D:/project/python/大洋测试/data/test.txt'\n",
    "    pretrain_model_path = 'D:/bert/hfl/chinese-roberta-wwm-ext'\n",
    "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "    model_final_save_path = 'D:/project/python/大洋测试/result/model/final'\n",
    "    model_save_path = 'D:/project/python/大洋测试/result/model'\n",
    "    # 训练参数\n",
    "    tokenizer = None\n",
    "    max_seq_len = 180\n",
    "    batch_size = 48\n",
    "    learning_rate = 1e-5\n",
    "    num_epochs = 10\n",
    "\n",
    "\n",
    "config = Config()      \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>label</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>布料一般，尺码合适，总体来说还算可以</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>这次购买太失败了，等了那么久才到，哎，希望京东可以自我检讨一下</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0</td>\n",
       "      <td>一点都不好用，第一次用，洗了头上感觉油油的，超级超级不喜欢，估计要丢了</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>裤子收到了 质量超级好 穿起来也很正式男朋友很喜欢</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0</td>\n",
       "      <td>送了根连接电视的接线，本本上却没有的接口，不知为什么！！！！！！！！ 发货时我选的圆通快递，...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  label                                               text\n",
       "0     1                                 布料一般，尺码合适，总体来说还算可以\n",
       "1     0                    这次购买太失败了，等了那么久才到，哎，希望京东可以自我检讨一下\n",
       "2     0                一点都不好用，第一次用，洗了头上感觉油油的，超级超级不喜欢，估计要丢了\n",
       "3     1                          裤子收到了 质量超级好 穿起来也很正式男朋友很喜欢\n",
       "4     0  送了根连接电视的接线，本本上却没有的接口，不知为什么！！！！！！！！ 发货时我选的圆通快递，..."
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def read_data(file):\n",
    "    data = []\n",
    "    with open(file, 'r') as f:\n",
    "        for line in f.readlines():\n",
    "            data.append(line.replace('\\n', '').split('\\t'))\n",
    "    return pd.DataFrame(data[1:], columns=['label', 'text'])\n",
    "train_df = read_data(config.train_data)\n",
    "train_df.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "180.0"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(config.pretrain_model_path)\n",
    "config.tokenizer = tokenizer\n",
    "train_df['len'] = train_df['text'].apply(lambda x: len(tokenizer.tokenize(x)))\n",
    "train_df['len'].quantile(0.95)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'input_ids': tensor([[ 101, 2791, 7313,  ...,    0,    0,    0],\n",
      "        [ 101,  794, 3025,  ...,    0,    0,    0],\n",
      "        [ 101,  671, 4692,  ...,    0,    0,    0],\n",
      "        ...,\n",
      "        [ 101,  679, 1962,  ...,    0,    0,    0],\n",
      "        [ 101, 3302, 1218,  ...,    0,    0,    0],\n",
      "        [ 101, 1678, 1678,  ...,    0,    0,    0]]), 'attention_mask': tensor([[1, 1, 1,  ..., 0, 0, 0],\n",
      "        [1, 1, 1,  ..., 0, 0, 0],\n",
      "        [1, 1, 1,  ..., 0, 0, 0],\n",
      "        ...,\n",
      "        [1, 1, 1,  ..., 0, 0, 0],\n",
      "        [1, 1, 1,  ..., 0, 0, 0],\n",
      "        [1, 1, 1,  ..., 0, 0, 0]]), 'labels': tensor([[1],\n",
      "        [0],\n",
      "        [0],\n",
      "        [1],\n",
      "        [1],\n",
      "        [1],\n",
      "        [0],\n",
      "        [1],\n",
      "        [0],\n",
      "        [1],\n",
      "        [0],\n",
      "        [1],\n",
      "        [0],\n",
      "        [0],\n",
      "        [0],\n",
      "        [1]])}\n"
     ]
    }
   ],
   "source": [
    "# 创建dataset\n",
    "class MyDataset(Dataset):\n",
    "    def __init__(self, config:Config, data:list, label:list=None):\n",
    "        self.config = config\n",
    "        self.data = data\n",
    "        self.label = label\n",
    "        self.max_seq_len = config.max_seq_len\n",
    "        self.tokenizer = config.tokenizer\n",
    "        self.mode = 'train' if label is not None else 'test'\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        text = self.data[idx]\n",
    "        inputs = self.tokenizer(\n",
    "            text,\n",
    "            max_length=self.max_seq_len,\n",
    "            return_attention_mask=True,\n",
    "            padding='max_length',\n",
    "            truncation=True\n",
    "        )\n",
    "        result = {\n",
    "            'input_ids': torch.tensor(inputs['input_ids'], dtype=torch.long),\n",
    "            'attention_mask': torch.tensor(inputs['attention_mask'], dtype=torch.long)\n",
    "        }\n",
    "        if self.mode == 'train':\n",
    "            result['labels'] = torch.tensor([int(self.label[idx])], dtype=torch.long)\n",
    "        return result\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "dataset = MyDataset(config,train_df['text'].tolist(),train_df['label'].to_list())\n",
    "dataloader = DataLoader(dataset, batch_size=16, shuffle=True)\n",
    "for idx, batch in enumerate(dataloader):\n",
    "    print(batch)\n",
    "    break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_data(data_path):\n",
    "    data = read_data(data_path)\n",
    "    mode = 'train' if 'label' in data.columns else 'test'\n",
    "    if mode == 'train':\n",
    "        X_train, X_test, y_train, y_test = train_test_split(data['text'].tolist(), data['label'].tolist(), test_size=0.2, random_state=42)\n",
    "        train_dataloader = DataLoader(MyDataset(config, X_train, y_train), batch_size=config.batch_size, shuffle=True)\n",
    "        test_dataloader = DataLoader(MyDataset(config, X_test, y_test), batch_size=config.batch_size, shuffle=False)\n",
    "        return train_dataloader, test_dataloader\n",
    "    else:\n",
    "        test_dataloader = DataLoader(MyDataset(config, data['text'].tolist(), None), batch_size=config.batch_size, shuffle=False)\n",
    "        return test_dataloader\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EMA():\n",
    "    def __init__(self, mu):\n",
    "        self.mu = mu\n",
    "        self.shadow = {}\n",
    " \n",
    "    def register(self, name, val):\n",
    "        self.shadow[name] = val.clone()\n",
    " \n",
    "    def get(self, name):\n",
    "        return self.shadow[name]\n",
    " \n",
    "    def update(self, name, x):\n",
    "        assert name in self.shadow\n",
    "        new_average = (1.0 - self.mu) * x + self.mu * self.shadow[name]\n",
    "        self.shadow[name] = new_average.clone()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at D:/bert/hfl/chinese-roberta-wwm-ext and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
      "0it [00:00, ?it/s]"
     ]
    }
   ],
   "source": [
    "def get_score(y_true, y_pred):\n",
    "    hit_count = sum([1 for i in range(len(y_true)) if y_pred[i] == y_true[i]])\n",
    "    probability = hit_count / len(y_pred)\n",
    "    return probability\n",
    "def val(model, test_dataloader:DataLoader):\n",
    "    model.eval()\n",
    "    total_acc, total_f1, test_num_batch = 0., 0., 0\n",
    "    total_y_pred, total_y_true = [], []\n",
    "    for batch in tqdm(test_dataloader):\n",
    "        batch_data = {k:v.to(config.device) for k, v in batch.items() }\n",
    "        with torch.no_grad():\n",
    "            outputs = model(**batch_data)\n",
    "            logits = outputs.logits\n",
    "            # 评估\n",
    "            y_pred = torch.argmax(logits, dim=1).cpu().detach().numpy()\n",
    "            y_true = batch_data['labels'].cpu().detach().numpy()\n",
    "            # 统计分数\n",
    "            acc = accuracy_score(y_true, y_pred)\n",
    "            f1 = f1_score(y_true, y_pred, average='weighted')\n",
    "            total_acc += acc\n",
    "            total_f1 += f1\n",
    "            test_num_batch += 1\n",
    "            total_y_pred.extend(y_pred)\n",
    "            total_y_true.extend(y_true)\n",
    "    # 计算平均分数\n",
    "    avg_acc = total_acc / test_num_batch\n",
    "    avg_f1 = total_f1 / test_num_batch\n",
    "    score = get_score(total_y_true, total_y_pred)\n",
    "    logger.info('-'*50)\n",
    "    logger.info(f\"val acc:{avg_acc}, val f1:{avg_f1}, val score:{score}\")\n",
    "    logger.info('-'*50)\n",
    "    return score\n",
    "\n",
    "\n",
    "def train(config:Config):\n",
    "    # 加载数据\n",
    "    train_dataloader, val_dataloader = load_data(config.train_data)\n",
    "    # 加载模型\n",
    "    model = AutoModelForSequenceClassification.from_pretrained(config.pretrain_model_path).to(config.device)\n",
    "    # 初始化优化器及梯度策略\n",
    "    optimizer = optim.AdamW(model.parameters(), lr=config.learning_rate)\n",
    "    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=config.num_epochs)\n",
    "\n",
    "    # ema\n",
    "    ema = EMA(0.999)\n",
    "    for name, param in model.named_parameters():\n",
    "        if param.requires_grad:\n",
    "            ema.register(name, param.data)\n",
    "\n",
    "    best_score = 0\n",
    "    # 遍历训练\n",
    "    for epoch in range(config.num_epochs):\n",
    "        model.train()\n",
    "        for idx, batch in tqdm(enumerate(train_dataloader)):\n",
    "            batch_data = {k:v.to(config.device) for k, v in batch.items() }\n",
    "            # 前向传播\n",
    "            outputs = model(**batch_data)\n",
    "            loss = outputs[0]\n",
    "            logits = outputs[1]\n",
    "            # 反向传播\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            scheduler.step()\n",
    "            optimizer.zero_grad()\n",
    "            # 评估\n",
    "            y_pred = torch.argmax(logits, dim=1).cpu().detach().numpy()\n",
    "            y_true = batch_data['labels'].cpu().detach().numpy()\n",
    "            acc = accuracy_score(y_true, y_pred)\n",
    "            f1 = f1_score(y_true, y_pred, average='weighted')\n",
    "            score = get_score(y_true, y_pred)\n",
    "            # 当准确率高于0.9时ema\n",
    "            if acc > 0.90:\n",
    "                for name, param in model.named_parameters():\n",
    "                    if param.requires_grad:\n",
    "                        ema.update(name, param.data)\n",
    "\n",
    "            if idx % 100 == 0:\n",
    "                logger.info(f\"epoch:{epoch}, idx:{idx}, loss:{loss.item()}, acc:{acc}, f1:{f1}, score:{score}\")\n",
    "        # 验证\n",
    "        val_score = val(model, val_dataloader)\n",
    "        if val_score > best_score:\n",
    "            best_score = val_score\n",
    "            model.save_pretrained(config.model_save_path)\n",
    "            tokenizer.save_pretrained(config.model_save_path)\n",
    "    \n",
    "    # 保存模型\n",
    "    model.save_pretrained(config.model_final_save_path)\n",
    "    tokenizer.save_pretrained(config.model_final_save_path)\n",
    "\n",
    "train(config)\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
